from .camvid import CamVid, get_camvid_loader

from .pascal_voc import get_voc_loader
from .cityscapes import Cityscapes, get_cityscapes_loader


def get_loader(args):
    if args.datasets == "camvid":
        return get_camvid_loader(root=args.data_path, batch_size=args.batch_size,
                                 train_crop_size=args.train_crop_size, test_crop_size=args.test_crop_size)

    elif args.datasets == "cityscapes":
        return get_cityscapes_loader()

    elif args.datasets == "voc":
        return get_voc_loader(root=args.data_path, batch_size=args.batch_size,
                              train_crop_size=args.train_crop_size, test_crop_size=args.test_crop_size,
                              crop_val=args.crop_val)
